// Copyright Epic Games, Inc. All Rights Reserved.

#include "fileremoteprojectstore.h"

#include <zencore/compress.h>
#include <zencore/filesystem.h>
#include <zencore/fmtutils.h>
#include <zencore/logging.h>
#include <zencore/timer.h>

namespace zen {

using namespace std::literals;

class LocalExportProjectStore : public RemoteProjectStore
{
public:
	LocalExportProjectStore(std::string_view			 Name,
							std::string_view			 OptionalBaseName,
							const std::filesystem::path& FolderPath,
							bool						 ForceDisableBlocks,
							bool						 ForceEnableTempBlocks)
	: m_Name(Name)
	, m_OptionalBaseName(OptionalBaseName)
	, m_OutputPath(FolderPath)
	{
		if (ForceDisableBlocks)
		{
			m_EnableBlocks = false;
		}
		if (ForceEnableTempBlocks)
		{
			m_UseTempBlocks = true;
		}
	}

	virtual RemoteStoreInfo GetInfo() const override
	{
		return {
			.CreateBlocks	   = m_EnableBlocks,
			.UseTempBlockFiles = m_UseTempBlocks,
			.AllowChunking	   = true,
			.ContainerName	   = m_Name,
			.BaseContainerName = m_OptionalBaseName,
			.Description =
				fmt::format("[file] {}/{}{}{}"sv, m_OutputPath, m_Name, m_OptionalBaseName.empty() ? "" : " Base: ", m_OptionalBaseName)};
	}

	virtual Stats GetStats() const override
	{
		return {.m_SentBytes		 = m_SentBytes.load(),
				.m_ReceivedBytes	 = m_ReceivedBytes.load(),
				.m_RequestTimeNS	 = m_RequestTimeNS.load(),
				.m_RequestCount		 = m_RequestCount.load(),
				.m_PeakSentBytes	 = m_PeakSentBytes.load(),
				.m_PeakReceivedBytes = m_PeakReceivedBytes.load(),
				.m_PeakBytesPerSec	 = m_PeakBytesPerSec.load()};
	}

	virtual SaveResult SaveContainer(const IoBuffer& Payload) override
	{
		Stopwatch  Timer;
		SaveResult Result;

		{
			CbObject ContainerObject = LoadCompactBinaryObject(Payload);

			ContainerObject.IterateAttachments([&](CbFieldView FieldView) {
				IoHash				  AttachmentHash = FieldView.AsBinaryAttachment();
				std::filesystem::path AttachmentPath = GetAttachmentPath(AttachmentHash);
				if (!std::filesystem::exists(AttachmentPath))
				{
					Result.Needs.insert(AttachmentHash);
				}
			});
		}

		std::filesystem::path ContainerPath = m_OutputPath;
		ContainerPath.append(m_Name);

		try
		{
			CreateDirectories(m_OutputPath);
			BasicFile ContainerFile;
			ContainerFile.Open(ContainerPath, BasicFile::Mode::kTruncate);
			std::error_code Ec;
			ContainerFile.WriteAll(Payload, Ec);
			if (Ec)
			{
				throw std::system_error(Ec, Ec.message());
			}
			Result.RawHash = IoHash::HashBuffer(Payload);
		}
		catch (const std::exception& Ex)
		{
			Result.ErrorCode = gsl::narrow<int32_t>(HttpResponseCode::InternalServerError);
			Result.Reason	 = fmt::format("Failed saving oplog container to '{}'. Reason: {}", ContainerPath, Ex.what());
		}
		AddStats(Payload.GetSize(), 0, Timer.GetElapsedTimeUs() * 1000);
		Result.ElapsedSeconds = Timer.GetElapsedTimeMs() / 1000.0;
		return Result;
	}

	virtual SaveAttachmentResult SaveAttachment(const CompositeBuffer& Payload, const IoHash& RawHash) override
	{
		Stopwatch			  Timer;
		SaveAttachmentResult  Result;
		std::filesystem::path ChunkPath = GetAttachmentPath(RawHash);
		if (!std::filesystem::exists(ChunkPath))
		{
			try
			{
				CreateDirectories(ChunkPath.parent_path());

				BasicFile ChunkFile;
				ChunkFile.Open(ChunkPath, BasicFile::Mode::kTruncate);
				size_t Offset = 0;
				for (const SharedBuffer& Segment : Payload.GetSegments())
				{
					ChunkFile.Write(Segment.GetView(), Offset);
					Offset += Segment.GetSize();
				}
			}
			catch (const std::exception& Ex)
			{
				Result.ErrorCode = gsl::narrow<int32_t>(HttpResponseCode::InternalServerError);
				Result.Reason	 = fmt::format("Failed saving oplog attachment to '{}'. Reason: {}", ChunkPath, Ex.what());
			}
		}
		AddStats(Payload.GetSize(), 0, Timer.GetElapsedTimeUs() * 1000);
		Result.ElapsedSeconds = Timer.GetElapsedTimeMs() / 1000.0;
		return Result;
	}

	virtual SaveAttachmentsResult SaveAttachments(const std::vector<SharedBuffer>& Chunks) override
	{
		Stopwatch Timer;

		for (const SharedBuffer& Chunk : Chunks)
		{
			CompressedBuffer	 Compressed	 = CompressedBuffer::FromCompressedNoValidate(Chunk.AsIoBuffer());
			SaveAttachmentResult ChunkResult = SaveAttachment(Compressed.GetCompressed(), Compressed.DecodeRawHash());
			if (ChunkResult.ErrorCode)
			{
				ChunkResult.ElapsedSeconds = Timer.GetElapsedTimeMs() / 1000.0;
				return SaveAttachmentsResult{ChunkResult};
			}
		}
		SaveAttachmentsResult Result;
		Result.ElapsedSeconds = Timer.GetElapsedTimeMs() / 1000.0;
		return Result;
	}

	virtual FinalizeResult FinalizeContainer(const IoHash&) override { return {}; }

	virtual LoadContainerResult LoadContainer() override { return LoadContainer(m_Name); }
	virtual LoadContainerResult LoadBaseContainer() override
	{
		if (m_OptionalBaseName.empty())
		{
			return LoadContainerResult{{.ErrorCode = static_cast<int>(HttpResponseCode::NoContent)}};
		}
		return LoadContainer(m_OptionalBaseName);
	}

	virtual HasAttachmentsResult HasAttachments(const std::span<IoHash> RawHashes) override
	{
		Stopwatch			 Timer;
		HasAttachmentsResult Result;
		for (const IoHash& RawHash : RawHashes)
		{
			std::filesystem::path ChunkPath = GetAttachmentPath(RawHash);
			if (!std::filesystem::is_regular_file(ChunkPath))
			{
				Result.Needs.insert(RawHash);
			}
		}
		AddStats(0, 0, Timer.GetElapsedTimeUs() * 1000);
		Result.ElapsedSeconds = Timer.GetElapsedTimeMs() / 1000.0;
		return Result;
	}

	virtual LoadAttachmentResult LoadAttachment(const IoHash& RawHash) override
	{
		Stopwatch			  Timer;
		LoadAttachmentResult  Result;
		std::filesystem::path ChunkPath = GetAttachmentPath(RawHash);
		if (!std::filesystem::is_regular_file(ChunkPath))
		{
			Result.ErrorCode = gsl::narrow<int>(HttpResponseCode::NotFound);
			Result.Reason = fmt::format("Failed loading oplog attachment from '{}'. Reason: 'The file does not exist'", ChunkPath.string());
			Result.ElapsedSeconds = Timer.GetElapsedTimeMs() / 1000.0;
			return Result;
		}
		{
			BasicFile ChunkFile;
			ChunkFile.Open(ChunkPath, BasicFile::Mode::kRead);
			Result.Bytes = ChunkFile.ReadAll();
		}
		AddStats(0, Result.Bytes.GetSize(), Timer.GetElapsedTimeUs() * 1000);
		Result.ElapsedSeconds = Timer.GetElapsedTimeMs() / 1000.0;
		return Result;
	}

	virtual LoadAttachmentsResult LoadAttachments(const std::vector<IoHash>& RawHashes) override
	{
		Stopwatch			  Timer;
		LoadAttachmentsResult Result;
		for (const IoHash& Hash : RawHashes)
		{
			LoadAttachmentResult ChunkResult = LoadAttachment(Hash);
			if (ChunkResult.ErrorCode)
			{
				ChunkResult.ElapsedSeconds = Timer.GetElapsedTimeMs() / 1000.0;
				return LoadAttachmentsResult{ChunkResult};
			}
			ZEN_DEBUG("Loaded attachment in {}", NiceTimeSpanMs(static_cast<uint64_t>(ChunkResult.ElapsedSeconds * 1000)));
			Result.Chunks.emplace_back(
				std::pair<IoHash, CompressedBuffer>{Hash, CompressedBuffer::FromCompressedNoValidate(std::move(ChunkResult.Bytes))});
		}
		return Result;
	}

private:
	LoadContainerResult LoadContainer(const std::string& Name)
	{
		Stopwatch			  Timer;
		LoadContainerResult	  Result;
		std::filesystem::path SourcePath = m_OutputPath;
		SourcePath.append(Name);
		if (!std::filesystem::is_regular_file(SourcePath))
		{
			Result.ErrorCode = gsl::narrow<int>(HttpResponseCode::NotFound);
			Result.Reason = fmt::format("Failed loading oplog container from '{}'. Reason: 'The file does not exist'", SourcePath.string());
			Result.ElapsedSeconds = Timer.GetElapsedTimeMs() / 1000.0;
			return Result;
		}
		IoBuffer ContainerPayload;
		{
			BasicFile ContainerFile;
			ContainerFile.Open(SourcePath, BasicFile::Mode::kRead);
			ContainerPayload = ContainerFile.ReadAll();
		}
		AddStats(0, ContainerPayload.GetSize(), Timer.GetElapsedTimeUs() * 1000);
		Result.ContainerObject = LoadCompactBinaryObject(ContainerPayload);
		if (!Result.ContainerObject)
		{
			Result.ErrorCode	  = gsl::narrow<int32_t>(HttpResponseCode::InternalServerError);
			Result.Reason		  = fmt::format("The file {} is not formatted as a compact binary object", SourcePath.string());
			Result.ElapsedSeconds = Timer.GetElapsedTimeMs() / 1000.0;
			return Result;
		}
		Result.ElapsedSeconds = Timer.GetElapsedTimeMs() / 1000.0;
		return Result;
	}

	std::filesystem::path GetAttachmentPath(const IoHash& RawHash) const
	{
		ExtendablePathBuilder<128> ShardedPath;
		ShardedPath.Append(m_OutputPath.c_str());
		ExtendableStringBuilder<64> HashString;
		RawHash.ToHexString(HashString);
		const char* str = HashString.c_str();
		ShardedPath.AppendSeparator();
		ShardedPath.AppendAsciiRange(str, str + 3);

		ShardedPath.AppendSeparator();
		ShardedPath.AppendAsciiRange(str + 3, str + 5);

		ShardedPath.AppendSeparator();
		ShardedPath.AppendAsciiRange(str + 5, str + 40);

		return ShardedPath.ToPath();
	}

	void AddStats(uint64_t UploadedBytes, uint64_t DownloadedBytes, uint64_t ElapsedNS)
	{
		m_SentBytes.fetch_add(UploadedBytes);
		m_ReceivedBytes.fetch_add(DownloadedBytes);
		m_RequestTimeNS.fetch_add(ElapsedNS);
		SetAtomicMax(m_PeakSentBytes, UploadedBytes);
		SetAtomicMax(m_PeakReceivedBytes, DownloadedBytes);
		if (ElapsedNS > 0)
		{
			uint64_t BytesPerSec = (gsl::narrow<uint64_t>(UploadedBytes + DownloadedBytes) * 1000000) / ElapsedNS;
			SetAtomicMax(m_PeakBytesPerSec, BytesPerSec);
		}

		m_RequestCount.fetch_add(1);
	}

	const std::string			m_Name;
	const std::string			m_OptionalBaseName;
	const std::filesystem::path m_OutputPath;
	bool						m_EnableBlocks	= true;
	bool						m_UseTempBlocks = false;

	std::atomic_uint64_t m_SentBytes;
	std::atomic_uint64_t m_ReceivedBytes;
	std::atomic_uint64_t m_RequestTimeNS;
	std::atomic_uint64_t m_RequestCount;
	std::atomic_uint64_t m_PeakSentBytes;
	std::atomic_uint64_t m_PeakReceivedBytes;
	std::atomic_uint64_t m_PeakBytesPerSec;
};

std::shared_ptr<RemoteProjectStore>
CreateFileRemoteStore(const FileRemoteStoreOptions& Options)
{
	std::shared_ptr<RemoteProjectStore> RemoteStore = std::make_shared<LocalExportProjectStore>(Options.Name,
																								Options.OptionalBaseName,
																								std::filesystem::path(Options.FolderPath),
																								Options.ForceDisableBlocks,
																								Options.ForceEnableTempBlocks);
	return RemoteStore;
}

}  // namespace zen
